import cv2
import numpy as np
import mediapipe as mp
from typing import List, Tuple, Generator
import torch
from pathlib import Path
import os
import logging
import insightface
from insightface.app import FaceAnalysis
import urllib.request
import time

logger = logging.getLogger(__name__)

class FaceSwapper:
    def __init__(self):
        """Initialize the face swapper with required models"""
        self.app = FaceAnalysis(name='buffalo_l')
        self.app.prepare(ctx_id=0 if torch.cuda.is_available() else -1, det_size=(640, 640))

        # Download and load the face swapping model
        model_path = self._download_model()
        self.swapper = insightface.model_zoo.get_model(model_path)

        # Load the source face (virtual face to swap with)
        self.source_face = self._load_source_face()
        self.source_face_embedding = None
        if self.source_face is not None:
            faces = self.app.get(self.source_face)
            if faces:
                self.source_face_embedding = faces[0]
                logger.info("Source face loaded successfully")
            else:
                logger.error("No face detected in source image")

    def _download_model(self) -> str:
        """Download the face swapping model if not exists"""
        model_path = Path("models/inswapper_128.onnx")
        if not model_path.exists():
            logger.info("Downloading face swapping model...")
            os.makedirs(model_path.parent, exist_ok=True)

            # List of mirror URLs
            urls = [
                "https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128.onnx",
            ]

            for url in urls:
                try:
                    logger.info(f"Trying to download from: {url}")
                    urllib.request.urlretrieve(url, model_path)
                    logger.info("Model downloaded successfully")
                    break
                except Exception as e:
                    logger.warning(f"Failed to download from {url}: {str(e)}")
                    continue

            if not model_path.exists():
                logger.error("Failed to download model from all sources")
                logger.info("Please manually download the model from one of these URLs and place it in models/inswapper_128.onnx")
                raise RuntimeError("Failed to download face swapping model")

        return str(model_path)

    def _load_source_face(self) -> np.ndarray:
        """Load the source face image"""
        # You should replace this with your own virtual face image
        source_path = Path("models/source_face.jpg")
        if not source_path.exists():
            logger.warning("Source face image not found. Please provide a source face image at models/source_face.jpg")
            return None
        return cv2.imread(str(source_path))

    def process_frames(self, frames_generator: Generator[List[np.ndarray], None, None]) -> Generator[List[np.ndarray], None, None]:
        """
        Process frames in batches, replacing faces with virtual ones

        Args:
            frames_generator: Generator yielding batches of frames

        Yields:
            Batches of processed frames
        """
        if self.source_face_embedding is None:
            logger.error("No source face available for swapping")
            for frames_batch in frames_generator:
                yield frames_batch
            return

        for frames_batch in frames_generator:
            processed_frames = []

            for frame in frames_batch:
                try:
                    # Detect faces in the frame
                    faces = self.app.get(frame)

                    if faces:
                        # Process each face in the frame
                        result = frame.copy()
                        for face in faces:
                            try:
                                result = self.swapper.get(
                                    result,
                                    face,
                                    self.source_face_embedding,
                                    paste_back=True
                                )
                            except Exception as e:
                                logger.error(f"Error swapping face: {str(e)}")
                        processed_frames.append(result)
                    else:
                        # If no face detected, keep original frame
                        processed_frames.append(frame)

                except Exception as e:
                    logger.error(f"Error processing frame: {str(e)}")
                    processed_frames.append(frame)  # Keep original frame on error

            yield processed_frames